from typing import Any, List, Dict
from calc_manager.base.plugins import PluginObject
from psycopg2 import pool
from dataclasses import dataclass
from enum import Enum
import numpy as np
import os


class ResultType(Enum):
    STRESS = 'stress'
    STRAIN = 'strain'
    FORCE = 'force'


@dataclass
class OP2ElementResultAttr:
    is_msc: bool
    is_nasa95: bool


@dataclass
class OP2ElementResultSpec:
    element_type: str
    result_type: str
    header: List[str]
    attrs: OP2ElementResultAttr
    version: str


@dataclass
class OP2Attr:
    subcases: List[int]
    result_types: List[str]
    element_types: List[str]


@dataclass
class OP2Spec:
    id: int
    file_name: str
    fp: str
    attrs: OP2Attr
    element_result_count: Dict[str, Any]


@dataclass
class OP2Data:
    eid: int
    element_type: str
    subcase: int
    element_node: int
    location: int
    result_type: str
    header: List[str]
    data: List[float]


@dataclass
class OP2ResultData:
    subcase: int
    element_node: int
    location: int
    data: List[float]


@dataclass
class OP2Result:
    eid: int
    element_type: str
    data: List[OP2ResultData]


@dataclass
class ShellElementEnvelope:
    eid: int
    result_type: str
    location: str
    header: List[str]
    max: List[float]
    min: List[float]


@dataclass
class ResultData:
    data: List[List[int]]
    element_node: List[List[int]]
    header: List[str]


class OP2Plugin(PluginObject):
    def __init__(self):
        self.db_pool = pool.SimpleConnectionPool(
            minconn=1,
            maxconn=3,
            host=os.environ.get("DB_HOST"),
            port=int(os.environ.get("DB_PORT")),
            dbname=os.environ.get("DB_DBNAME"),
            user=os.environ.get("DB_USER"),
            password=os.environ.get("DB_PASSWORD")
        )
        self.element_result_cache = {}
        print("op2 plugin init......")

    def get_element_result_spec(self, element_type: str, result_type: str) -> OP2ElementResultSpec:
        """
        根据element_type、result_type获取属性信息（header、 attrs）
        :param element_type:
        :param result_type:
        :return:
        """
        key = f'{element_type}_{result_type}'
        if key in self.element_result_cache.keys():
            return self.element_result_cache[key]
        query_sql = f'select * from "op2_element_result_attr" where element_type = %s and result_type = %s'
        conn = self.db_pool.getconn()
        try:
            with conn.cursor() as cur:
                cur.execute(query_sql, (element_type, result_type))
                res = cur.fetchone()
        except Exception as e:
            print(f"Error: {e}")
        finally:
            self.db_pool.putconn(conn)
        if res:
            ret = OP2ElementResultSpec(
                element_type=element_type,
                result_type=result_type,
                header=[var.strip() for var in res[2].split(",")],
                attrs=OP2ElementResultAttr(**res[3]),
                version=res[4]
            )
            self.element_result_cache[key] = ret
            return ret

    def get_op2_spec(self, op2_fp) -> OP2Spec:
        """
        根据op2文件指纹获取op2文件元信息
        :param op2_fp:
        :return:
        """
        query_sql = f'select * from op2_file_resource where fp = %s'
        conn = self.db_pool.getconn()
        try:
            with conn.cursor() as cur:
                cur.execute(query_sql, (op2_fp,))
                res = cur.fetchone()
        except Exception as e:
            print(f"Error: {e}")
        finally:
            self.db_pool.putconn(conn)
        if res:
            return OP2Spec(
                id=res[0],
                file_name=res[1],
                fp=res[2],
                attrs=OP2Attr(**res[3]),
                element_result_count=res[4]
            )

    def get_result(self, op2_fp: str, eid: int, result_type: str,
                   subcase: int = None, element_node: int = None, location: int = None) -> List[OP2Data]:
        """
        根据op2文件名、eid、结果类型获取单元结果
        :param op2_fp: op2文件指纹
        :param eid: eid
        :param result_type: stress: 应力  strain：应变
        :param subcase: subcase
        :param element_node: 单元节点编号
        :param location: z1 or z2
        :return:
        """
        ret = []
        query_sql = f'select * from "op2_{op2_fp}_{result_type}" where eid = {eid}'
        conn = self.db_pool.getconn()
        try:
            with conn.cursor() as cur:
                cur.execute(query_sql)
                res = cur.fetchone()
        except Exception as e:
            print(f"Error: {e}")
        finally:
            self.db_pool.putconn(conn)
        if res is None:
            return ret
        stress_datas = res[2]
        for stress_data in stress_datas:
            if subcase and stress_data['subcase'] != subcase:
                continue
            if element_node and stress_data['element_node'] != element_node:
                continue
            if location and stress_data['location'] != location:
                continue
            ret.append(OP2Data(eid=res[0],
                               element_type=res[1],
                               subcase=stress_data['subcase'],
                               element_node=stress_data['element_node'],
                               location=stress_data['location'],
                               result_type=result_type,
                               header=self.get_element_result_spec(res[1], result_type).header,
                               data=stress_data['data']))
        return ret

    def get_results(self, op2_fp: str, eids: List[int], result_type: str) -> List[OP2Result]:
        """
        根据op2文件名、eid、结果类型获取单元结果
        :return:
        """
        ret = []
        query_sql = f'select * from "op2_{op2_fp}_{result_type}" where eid = ANY(%s)'
        conn = self.db_pool.getconn()
        try:
            with conn.cursor() as cur:
                cur.execute(query_sql, (eids,))
                res = cur.fetchall()
        except Exception as e:
            print(f"Error: {e}")
        finally:
            self.db_pool.putconn(conn)
        if res is None:
            return ret
        for row in res:
            ret.append(OP2Result(eid=row[0],
                                 element_type=row[1],
                                 data=row[2]))
        return ret

    def get_result(self, op2_fp: str, eids: List[int], subcases: List[int], element_type: str, result_type: str) -> Dict[int, ResultData]:
        """
        根据subcase分类结果
        :param op2_fp:
        :param eids:
        :param subcases:
        :param element_type:
        :param result_type:
        :return: subcase: ResultData
        """
        ret = {}
        query_sql = f'select * from "op2_{op2_fp}_{result_type}" where element_type = %s and eid = ANY(%s)'
        conn = self.db_pool.getconn()
        try:
            with conn.cursor() as cur:
                cur.execute(query_sql, (element_type, eids))
                res = cur.fetchall()
        except Exception as e:
            print(f"Error: {e}")
        finally:
            self.db_pool.putconn(conn)
        if res is None:
            return ret
        for subcase in subcases:
            data_list = []
            element_node_list = []
            for row in res:
                for node_data in row[2]:
                    if node_data['subcase'] != subcase:
                        continue
                    data_list.append(node_data['data'])
                    element_node_list.append([row[0], node_data['element_node']])
            ret[subcase] = ResultData(
                data=data_list,
                element_node=element_node_list,
                header=self.get_element_result_spec(element_type, result_type).header,
            )
        return ret

    def get_shell_element_envelope(self, op2_fp: str, eids: List[int], subcases: List[int] = None,
                                   result_types: List[str] = None, location: str = "all") -> List[ShellElementEnvelope]:
        """
        计算壳单元包络
        :param op2_fp: op2文件指纹
        :param eids: eid list
        :param subcases: subcase list
        :param result_types: ["stress", "strain"]
        :param location: "z1", "z2" or "all"
        :return:
        """
        if result_types is None:
            result_types = [ResultType.STRESS.value, ResultType.STRAIN.value]
        ret = []
        for result_type in result_types:
            query_sql = f'select * from "op2_{op2_fp}_{result_type}" where element_type IN %s and eid = ANY(%s)'
            conn = self.db_pool.getconn()
            res = None
            try:
                with conn.cursor() as cur:
                    cur.execute(query_sql, (('ctria3', 'cquad4'), eids))
                    res = cur.fetchall()
            except Exception as e:
                print(f"Error: {e}")
            finally:
                self.db_pool.putconn(conn)

            if res is None:
                continue
            for row in res:
                datas = []
                for node_data in row[2]:
                    if subcases and node_data['subcase'] not in subcases:
                        continue
                    if location == "all" or location == node_data['location']:
                        datas.append(node_data['data'])
                if len(datas) == 0:
                    continue
                matrix = np.array(datas)
                max_values = np.max(matrix, axis=0)
                min_values = np.min(matrix, axis=0)
                ret.append(ShellElementEnvelope(
                    eid=row[0],
                    result_type=result_type,
                    location=location,
                    header=self.get_element_result_spec('ctria3', result_type).header,
                    max=list(max_values),
                    min=list(min_values)
                ))
        return ret
